{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Week8 bonus descriptions\n",
    "\n",
    "Here are some cool mini-projects you can try to dive deeper into the topic.\n",
    "\n",
    "## More metrics: BLEU (5+ pts)\n",
    "\n",
    "Pick BLEU or any other relevant metric, e.g. BLEU (e.g. from `nltk.bleu_score`).\n",
    "* Train model to maximize BLEU directly\n",
    "* How does levenshtein behave when maximizing BLEU and vice versa?\n",
    "* Compare this with how they behave when optimizing likelihood. \n",
    "\n",
    "(use default parameters for bleu: 4-gram, uniform weights)\n",
    "\n",
    "## Actor-critic (5+++ pts)\n",
    "\n",
    "While self-critical training provides a large reduction of gradient variance, it has a few drawbacks:\n",
    "- It requires a lot of additional computation during training\n",
    "- It doesn't adjust V(s) between decoder steps. (one value per sequence)\n",
    "\n",
    "There's a more general way of doing the same thing: learned baselines, also known as __advantage actor-critic__.\n",
    "\n",
    "There are two main ways to apply that:\n",
    "- __naive way__: compute V(s) once per training example.\n",
    "  - This only requires additional 1-unit linear dense layer that grows out of encoder, estimating V(s)\n",
    "  - (implement this to get main points)\n",
    "- __every step__: compute V(s) on each decoder step\n",
    "  - Again it's just an 1-unit dense layer (no nonlinearity), but this time it's inside decoder recurrence.\n",
    "  - (+3 pts additional for this guy)\n",
    "\n",
    "In both cases, you should train V(s) to minimize squared error $(V(s) - R(s,a))^2$ with R being actual levenshtein.\n",
    "You can then use $ A(s,a) = (R(s,a) - const(V(s))) $ for policy gradient.\n",
    "\n",
    "There's also one particularly interesting approach (+5 additional pts):\n",
    "- __combining SCST and actor-critic__:\n",
    "  - compute baseline $V(s)$ via self-critical sequence training (just like in main assignment)\n",
    "  - learn correction $ C(s,a_{:t}) = R(s,a) - V(s) $ by minimizing $(R(s,a) - V(s) - C(s,a_{:t}))^2 $\n",
    "  - use $ A(s,a_{:t}) = R(s,a) - V(s) - const(C(s,a_{:t})) $\n",
    "\n",
    "\n",
    "\n",
    "## Implement attention (5+++ pts)\n",
    "\n",
    "Some seq2seq tasks can benefit from the attention mechanism. In addition to taking the _last_ time-step of encoder hidden state, we can allow decoder to peek on any time-step of his choice.\n",
    "\n",
    "![img](https://xiandong79.github.io/downloads/nmt-model-fast.gif)\n",
    "\n",
    "\n",
    "#### Recommended steps:\n",
    "__1)__ Modify encoder-decoder\n",
    "\n",
    "Learn to feed the entire encoder into the decoder. You can do so by sending encoder rnn sequences directly into decoder (make sure there's no `only_return_final=True` for encoder rnn layer).\n",
    "\n",
    "```\n",
    "class encoder:\n",
    "    ...\n",
    "    enc_sequences, (h, c) = self.lstm(x)\n",
    "    ...\n",
    "    \n",
    "class decoder: \n",
    "    ...\n",
    "    attention_applied = self.attn_layer(enc_sequences)\n",
    "    h, c = self.lstm_decoder(prev_emb, (attention_applied, c))\n",
    "    ...\n",
    "    \n",
    "```\n",
    "    \n",
    "\n",
    "For starters, you can take it's last tick (via SliceLayer) inside the decoder step and feed it as input to make sure it works.\n",
    "\n",
    "__2)__ Implement attention mechanism\n",
    "\n",
    "Next thing we'll need is to implement the math of attention.\n",
    "\n",
    "The simplest way to do so is to write a special layer. We gave you a prototype and some tests below.\n",
    "\n",
    "__3)__ Use attention inside decoder\n",
    "\n",
    "That's almost it! Now use `AttentionLayer` inside the decoder and feed it to back to lstm/gru/rnn (see code demo below).\n",
    "\n",
    "Train the full network just like you did before attention.\n",
    "\n",
    "__More points__ will be awwarded for comparing learning results of attention Vs no attention.\n",
    "\n",
    "__Bonus bonus:__ visualize attention vectors (>= +3 points)\n",
    "\n",
    "The best way to make sure your attention actually works is to visualize it.\n",
    "\n",
    "A simple way to do so is to obtain attention vectors from each tick (values __right after softmax__, not the layer outputs) and drawing those as images.\n",
    "\n",
    "#### step-by-step guide:\n",
    "\n",
    "- compute scores between $h_{e, j}^i$ and $h_{d}^i$ $\\forall j = 1, ... , \\text{len(enc_seq)}$, where i -- number of decoder step\n",
    "- apply softmax to scores and get weight for each vector\n",
    "- obtain attention vector using enc_seq and weights\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionLayer(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size):\n",
    "        super(self.__class__, self).__init__()\n",
    "        self.attn = nn.Linear(hidden_size * 2, output_size)\n",
    "\n",
    "    \n",
    "    def forward(self, enc_seq, decoder_state): \n",
    "        scores = # your code here\n",
    "        alphas = F.softmax(\n",
    "            # your code here\n",
    "        )\n",
    "        attn_combined = # your code here\n",
    "        return attn_combined"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# demo code\n",
    "batch_size = 32\n",
    "hidden_size = 256\n",
    "seq_len = 41\n",
    "dec_h_prev = torch.rand((batch_size, hidden_size))\n",
    "enc_sequences = torch.rand((batch_size, seq_len, hidden_size))\n",
    "\n",
    "attention = AttentionLayer(hidden_size, hidden_size)\n",
    "\n",
    "# sanity check\n",
    "demo_output = attention(enc_sequences, dec_h_prev)\n",
    "print('actual shape:', demo_output.shape)\n",
    "assert demo_output.shape == (32, 256)\n",
    "assert np.all(np.isfinite(demo_output.detach().cpu().numpy()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
